/*
 * Copyright (c) 2022 Xiaomi Coporation
 *
 * SPDX-License-Identifier: Apache-2.0
 */

/* LzmaUtil.c -- Test application for LZMA compression
   2018-07-04 : Igor Pavlov : Public domain */

#include "Precomp.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "7zFile.h"
#include "Alloc.h"
#include "LzmaDec.h"
#include "LzmaEnc.h"

typedef struct {
	unsigned char major;
	unsigned char minor;
	unsigned char revision;
	unsigned char padding;
	unsigned short developer_version;

	unsigned short p;
	unsigned int m;

	unsigned int magic;
} info_t;

static int PrintError(const char *message) {
	fprintf(stderr, "%s\n", message);
	return 1;
}

#define TLSR82_START_ADDR		(0)
#define TLSR82_TOTAL_SIZE		(512 * 1024)

#define TLSR82_BOOT_ADDR		TLSR82_START_ADDR
#define TLSR82_BOOT_SIZE		(20 * 1024)

#define TLSR82_BOOT_SETTING_ADDR	(TLSR82_BOOT_ADDR + TLSR82_BOOT_SIZE)
#define TLSR82_BOOT_SETTING_SIZE	(4 * 1024)

#define TLSR82_OTA_BANK_1_ADDR		(TLSR82_BOOT_SETTING_ADDR + TLSR82_BOOT_SETTING_SIZE)
#define TLSR82_OTA_BANK_1_SIZE		(104 * 1024)

#define TLSR82_APP_ADDR			(TLSR82_OTA_BANK_1_ADDR + TLSR82_OTA_BANK_1_SIZE)
#define TLSR82_APP_SIZE			(272 * 1024)

#define TLSR82_OTA_BANK_2_ADDR		(TLSR82_APP_ADDR + TLSR82_APP_SIZE)
#define TLSR82_OTA_BANK_2_SIZE		(92 * 1024)

#define TLSR82_NVS_ADDR			(TLSR82_OTA_BANK_2_ADDR + TLSR82_OTA_BANK_2_SIZE)
#define TLSR82_NVS_SIZE			(16 * 1024)

#define TLSR82_CERT_ADDR		(TLSR82_NVS_ADDR + TLSR82_NVS_SIZE)
#define TLSR82_CERT_SIZE		(4 * 1024)

#define ALIGN_NUM(val, num)		((val + num - 1) / num)

static void show_help(void)
{
	fprintf(stdout, "Version: 1.0.3\n");
	fprintf(stderr, "Usage: ./merge <path>/app.bin [<path>/boot.bin"
	        " <path>/bqb_test.bin <path>/factory.bin] \n");
}

static SRes Encode(ISeqOutStream *outStream, ISeqInStream *inStream,
                   UInt64 fileSize, char *rs) {
	CLzmaEncHandle enc;
	SRes res;
	CLzmaEncProps props;

	UNUSED_VAR(rs);

	enc = LzmaEnc_Create(&g_Alloc);
	if (enc == 0)
		return SZ_ERROR_MEM;

	LzmaEncProps_Init(&props);
	res = LzmaEnc_SetProps(enc, &props);

	if (res == SZ_OK) {
		Byte header[LZMA_PROPS_SIZE + 8];
		size_t headerSize = LZMA_PROPS_SIZE;
		int i;

		res = LzmaEnc_WriteProperties(enc, header, &headerSize);
		for (i = 0; i < 8; i++)
			header[headerSize++] = (Byte)(fileSize >> (8 * i));
		if (outStream->Write(outStream, header, headerSize) != headerSize)
			res = SZ_ERROR_WRITE;
		else {
			if (res == SZ_OK)
				res = LzmaEnc_Encode(enc, outStream, inStream, NULL, &g_Alloc, &g_Alloc);
		}
	}

	LzmaEnc_Destroy(enc, &g_Alloc, &g_Alloc);
	return res;
}

static int main2(const char *in, const char *out)
{
	CFileSeqInStream inStream;
	CFileOutStream outStream;
	int res;
	UInt64 fileSize;

	FileSeqInStream_CreateVTable(&inStream);
	File_Construct(&inStream.file);

	FileOutStream_CreateVTable(&outStream);
	File_Construct(&outStream.file);

	if (InFile_Open(&inStream.file, in) != 0)
		return PrintError("Can not open input file");

	if (OutFile_Open(&outStream.file, out) != 0)
		return PrintError("Can not open output file");

	File_GetLength(&inStream.file, &fileSize);
	res = Encode(&outStream.vt, &inStream.vt, fileSize, NULL);

	File_Close(&outStream.file);
	File_Close(&inStream.file);

	return res;
}

int main(int argc, char **argv) {
	char *boot_file, *fac_file, *app_file, *app_ota_file, *app_lzma_file, *outName;
	FILE *boot_File, *fac_File, *app_File, *app_ota_File, *app_lzma_File, *outFile;
	unsigned char *total_buf = NULL;
	size_t app_file_size = 0;
	size_t app_ota_file_size = 0;
	size_t app_lzma_file_size = 0;
	size_t app_ota_bank_1_size = 0;
	size_t app_ota_bank_2_size = 0;
	size_t boot_file_size = 0;
	size_t fac_file_size = 0;
	size_t total_len = 0;
	char path_lzma[128];
	char path_ota[128];
	char header[16] = {0};
	size_t path_len;
	info_t info;

	// Default arguments
	boot_file = NULL;
	fac_file = NULL;
	app_file = NULL;
	app_ota_file = NULL;
	app_lzma_file = NULL;
	outName = NULL;

	boot_File = NULL;
	fac_File = NULL;
	app_File = NULL;
	app_ota_File = NULL;
	app_lzma_File = NULL;
	outFile = NULL;

	if ((argc != 5 && argc != 2) || sizeof(info_t) != 16) {
		show_help();
		return 0;
	}

	app_file = argv[1];
	boot_file = argv[2];
	fac_file = argv[3];
	outName = argv[4];

	app_File = fopen(app_file, "rb");
	if (!app_File) {
		fprintf(stderr, "Unable to open app file \"%s\".\n", app_file);
		show_help();
		return 0;
	}

	fseek(app_File, 0, SEEK_END);
	app_file_size = (size_t)ftell(app_File);
	fseek(app_File, app_file_size - sizeof(info_t), SEEK_SET);

	if (app_file_size > TLSR82_APP_SIZE) {
		fprintf(stderr, "App file \"%s\" too bigger than (%ld > %dKB).\n", app_file,
		        app_file_size / 1024, TLSR82_APP_SIZE / 1024);
		goto _free;
	}

	if (fread(&info, 1, sizeof(info_t), app_File) != sizeof(info_t)) {
		goto _free;
	}

	fclose(app_File);
	app_File = NULL;

	path_len = strlen(app_file);
	memcpy(path_lzma + 1, app_file, path_len);
	path_len += 1;

	path_lzma[0] = '.';

	path_lzma[path_len + 0] = '.';
	path_lzma[path_len + 1] = 'l';
	path_lzma[path_len + 2] = 'z';
	path_lzma[path_len + 3] = 'm';
	path_lzma[path_len + 4] = 'a';
	path_lzma[path_len + 5] = 0;

	if (main2(app_file, path_lzma)) {
		fprintf(stderr, "Unable to create lzma file \"%s\".\n", path_lzma);
		return 0;
	}

	total_buf = (unsigned char *)malloc(TLSR82_TOTAL_SIZE);
	if (!total_buf) {
		return 0;
	}

	app_lzma_file = path_lzma;

	app_lzma_File = fopen(app_lzma_file, "rb");
	if (!app_lzma_File) {
		fprintf(stderr, "Unable to open app lzma file \"%s\".\n", app_lzma_file);
		show_help();
		goto _free;
	}

	fseek(app_lzma_File, 0, SEEK_END);
	app_lzma_file_size = (size_t)ftell(app_lzma_File);
	fseek(app_lzma_File, 0, SEEK_SET);

	if (fread(total_buf, 1, app_lzma_file_size, app_lzma_File) !=
		  app_lzma_file_size) {
		goto _free;
	}

	fclose(app_lzma_File);
	app_lzma_File = NULL;

	snprintf(path_ota, sizeof(path_ota), "tlsr_ota_pid_%d_ver_%d_%d_%d_%04d.bin",
		 info.p, info.major, info.minor, info.revision, info.developer_version);

	app_ota_file = path_ota;

	app_ota_File = fopen(app_ota_file, "wb+");
	if (!app_ota_File) {
		fprintf(stderr, "Unable to open app ota file \"%s\".\n", app_ota_file);
		show_help();
		goto _free;
	}

	snprintf(header, sizeof(header), "miot %08x", (unsigned int)app_lzma_file_size);

	if (fwrite(header, 1, sizeof(header), app_ota_File) != sizeof(header)) {
		fprintf(stderr, "Error writing header to ota file.\n");
	}

	if (fwrite(total_buf, 1, app_lzma_file_size, app_ota_File) != app_lzma_file_size) {
		fprintf(stderr, "Error writing raw to ota file.\n");
	}

	fclose(app_ota_File);
	app_ota_File = NULL;

	app_ota_File = fopen(app_ota_file, "rb");
	if (!app_ota_File) {
		fprintf(stderr, "Unable to open app ota file \"%s\".\n", app_ota_file);
		show_help();
		goto _free;
	}

	fseek(app_ota_File, 0, SEEK_END);
	app_ota_file_size = (size_t)ftell(app_ota_File);
	fseek(app_ota_File, 0, SEEK_SET);

	printf("APP/OTA(compressed)\t\t%ldKB/%ldKB \tPID %d Version %d.%d.%d_%04d\n",
	       ALIGN_NUM(app_file_size, 1024), ALIGN_NUM(app_ota_file_size, 1024),
	       info.p, info.major, info.minor, info.revision, info.developer_version);

	if (app_ota_file_size > (TLSR82_OTA_BANK_1_SIZE + TLSR82_OTA_BANK_2_SIZE)) {
		fprintf(stderr, "App ota file \"%s\" too bigger than (%ld > %dKB).\n",
		        app_ota_file, app_ota_file_size / 1024,
		        (TLSR82_OTA_BANK_1_SIZE + TLSR82_OTA_BANK_2_SIZE) / 1024);
		goto _free;
	}

	if (argc == 2) {
		goto _free;
	}

	(void)memset(total_buf, 0xff, TLSR82_TOTAL_SIZE);

	boot_File = fopen(boot_file, "rb");
	if (!boot_File) {
		fprintf(stderr, "Unable to open boot file \"%s\".\n", boot_file);
		goto _free;
	}

	fseek(boot_File, 0, SEEK_END);
	boot_file_size = (size_t)ftell(boot_File);
	fseek(boot_File, 0, SEEK_SET);

	printf("Boot size\t\t\t%ldKB\n", ALIGN_NUM(boot_file_size, 1024));

	if (boot_file_size > TLSR82_BOOT_SIZE) {
		fprintf(stderr, "Boot file \"%s\" too bigger than (%ld > %dKB).\n",
		        boot_file, boot_file_size / 1024, TLSR82_BOOT_SIZE / 1024);
		goto _free;
	}

	if (fread(total_buf, 1, boot_file_size, boot_File) != boot_file_size) {
		goto _free;
	}

	fclose(boot_File);
	boot_File = NULL;

	total_buf[8] = 0xFF;
	total_len = TLSR82_BOOT_SETTING_ADDR;
	(void)memcpy(total_buf + total_len, total_buf, TLSR82_BOOT_SETTING_SIZE);
	total_len += TLSR82_BOOT_SETTING_SIZE;

	app_ota_bank_1_size = (app_ota_file_size > TLSR82_OTA_BANK_1_SIZE)
	                       ? (TLSR82_OTA_BANK_1_SIZE)
	                       : app_ota_file_size;

	if (fread(total_buf + total_len, 1, app_ota_bank_1_size, app_ota_File) !=
	    app_ota_bank_1_size) {
		goto _free;
	}

	total_len = TLSR82_APP_ADDR;

	fac_File = fopen(fac_file, "rb");
	if (!fac_File) {
		fprintf(stderr, "Unable to open factory file \"%s\".\n", fac_file);
		show_help();
		goto _free;
	}

	fseek(fac_File, 0, SEEK_END);
	fac_file_size = (size_t)ftell(fac_File);
	fseek(fac_File, 0, SEEK_SET);

	printf("Factory size\t\t\t%ldKB\n", ALIGN_NUM(fac_file_size, 1024));

	if (fac_file_size > 100 * 1024) {
		fprintf(stderr, "Factory file \"%s\" too bigger than (%ld > 100KB).\n",
		        fac_file, fac_file_size / 1024);

		goto _free;
	}

	if (fread(total_buf + total_len, 1, fac_file_size, fac_File) !=
	    fac_file_size) {
		goto _free;
	}

	fclose(fac_File);
	fac_File = NULL;

	total_len += fac_file_size;

	memcpy(total_buf + total_len, &info, sizeof(info_t));

	total_len = ALIGN_NUM((total_len + sizeof(info_t)), 4096);
	total_len *= 4096;

	app_ota_bank_2_size = app_ota_file_size - app_ota_bank_1_size;
	if (fread(total_buf + total_len, 1, app_ota_bank_2_size, app_ota_File) !=
		  app_ota_bank_2_size) {
		goto _free;
	}

	fclose(app_ota_File);
	app_ota_File = NULL;

	total_len += app_ota_bank_2_size;

	// Compressed data is now in encBuf, write it...
	outFile = fopen(outName, "wb");
	if (!outFile) {
		fprintf(stderr, "Unable to open out file \"%s\".\n", outName);
		show_help();
		goto _free;
	}

	// Write header
	if (fwrite(total_buf, 1, total_len, outFile) != total_len) {
		fprintf(stderr, "Error writing to output file.\n");
	}

	fclose(outFile);

	printf("Total len\t\t\t%ldKB\n", ALIGN_NUM(total_len, 1024));

	printf("\n\n\nDone.\n\n\n");

_free:
	if (boot_File) {
		fclose(boot_File);
	}

	if (fac_File) {
		fclose(fac_File);
	}

	if (app_File) {
		fclose(app_File);
	}

	if (app_ota_File) {
		fclose(app_ota_File);
	}

	if (app_lzma_File) {
		fclose(app_lzma_File);
	}

	// Free memory
	if (total_buf) {
		free(total_buf);		
	}

	return 0;
}
